import logging
from django.utils import timezone
from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework_simplejwt.tokens import RefreshToken
from rest_framework_simplejwt.exceptions import InvalidToken
from .models import BlacklistedToken

logger = logging.getLogger(__name__)


class CustomJWTAuthentication(JWTAuthentication):
    """
    自定义JWT认证类，增加黑名单令牌检查功能
    """
    
    def authenticate(self, request):
        """
        验证JWT令牌，并检查是否在黑名单中
        """
        header = self.get_header(request)
        if header is None:
            return None

        raw_token = self.get_raw_token(header)
        if raw_token is None:
            return None

        # 检查令牌是否在黑名单中
        if BlacklistedToken.objects.filter(token=raw_token, expires_at__gt=timezone.now()).exists():
            logger.warning(f"令牌已在黑名单中: {raw_token}")
            raise InvalidToken("令牌已在黑名单中")

        validated_token = self.get_validated_token(raw_token)
        user = self.get_user(validated_token)
        
        # 添加认证日志
        logger.info(f"用户 {user.id} 使用令牌认证成功")
        return (user, validated_token)
    
    def authenticate_refresh_token(self, refresh_token):
        """
        验证刷新令牌，并检查是否在黑名单中
        """
        # 检查刷新令牌是否在黑名单中
        if BlacklistedToken.objects.filter(token=refresh_token, expires_at__gt=timezone.now()).exists():
            logger.warning(f"刷新令牌已在黑名单中: {refresh_token}")
            raise InvalidToken("刷新令牌已在黑名单中")
        
        # 验证刷新令牌
        return RefreshToken(refresh_token)